from scipy.io import wavfile
import os
import sys

current_dir = os.path.dirname(os.path.abspath(__file__))
my_own_dir = os.path.join(current_dir, 'my_own_')
sys.path.append(my_own_dir)


import numpy as np
import torch
import utils
import argparse
from scipy.io import wavfile
from text.symbols import symbols
from text import cleaned_text_to_sequence
from vits_pinyin import VITS_PinYin

def save_wav(wav, path, rate):
    wav *= 32767 / max(0.01, np.max(np.abs(wav))) * 0.6
    wavfile.write(path, rate, wav.astype(np.int16))
    
def gen_wave(item, wav_name = "temp.wav"):
    try:
        original_dir = os.getcwd()
        current_dir = os.path.dirname(os.path.abspath(__file__))
        y_own_dir = os.path.join(current_dir, 'my_own_')
        os.chdir(y_own_dir)
        out_path = os.path.join(y_own_dir, 'vits_infer_out')
        saved_wav_path = os.path.join(out_path, wav_name)
        config_path = './configs/bert_vits_student.json'
        model_path = './logs/bert_vits/vits_bert_student.pth'

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        tts_front = VITS_PinYin("./bert", device)
        hps = utils.get_hparams_from_file(config_path)
        net_g = utils.load_class(hps.train.eval_class)(
            len(symbols),
            hps.data.filter_length // 2 + 1,
            hps.train.segment_size // hps.data.hop_length,
            **hps.model)

        utils.load_model(model_path, net_g)
        net_g.eval()
        net_g.to(device)


        os.makedirs(out_path, exist_ok=True)



        phonemes, char_embeds = tts_front.chinese_to_phonemes(item)
        input_ids = cleaned_text_to_sequence(phonemes)
        with torch.no_grad():
            x_tst = torch.LongTensor(input_ids).unsqueeze(0).to(device)
            x_tst_lengths = torch.LongTensor([len(input_ids)]).to(device)
            x_tst_prosody = torch.FloatTensor(char_embeds).unsqueeze(0).to(device)
            audio = net_g.infer(x_tst, x_tst_lengths, x_tst_prosody, noise_scale=0.5,
                                length_scale=1)[0][0, 0].data.cpu().float().numpy()

        os.chdir(original_dir)
        save_wav(audio, saved_wav_path, hps.data.sampling_rate)
        
    finally:
        os.chdir(original_dir)
    return saved_wav_path

